import React, { useEffect, useRef } from 'react';

interface ConfusionMatrixProps {
  matrix: number[][];
}

const ConfusionMatrix: React.FC<ConfusionMatrixProps> = ({ matrix }) => {
  const canvasRef = useRef<HTMLCanvasElement>(null);
  
  useEffect(() => {
    if (!canvasRef.current || !matrix || matrix.length === 0) return;
    
    const canvas = canvasRef.current;
    const ctx = canvas.getContext('2d');
    if (!ctx) return;
    
    // Set canvas dimensions
    const cellSize = 40;
    const padding = 50;
    const numClasses = matrix.length;
    const width = numClasses * cellSize + padding * 2;
    const height = numClasses * cellSize + padding * 2;
    
    canvas.width = width;
    canvas.height = height;
    
    // Clear canvas
    ctx.clearRect(0, 0, width, height);
    
    // Find max value for color scaling
    let maxVal = 0;
    for (let i = 0; i < numClasses; i++) {
      for (let j = 0; j < numClasses; j++) {
        maxVal = Math.max(maxVal, matrix[i][j]);
      }
    }
    
    // Draw matrix
    for (let i = 0; i < numClasses; i++) {
      for (let j = 0; j < numClasses; j++) {
        const value = matrix[i][j];
        const normalized = maxVal > 0 ? value / maxVal : 0;
        
        // Use a blue gradient
        const intensity = Math.floor(255 * normalized);
        const color = `rgb(${255 - intensity}, ${255 - Math.floor(intensity * 0.5)}, 255)`;
        
        ctx.fillStyle = color;
        ctx.fillRect(
          j * cellSize + padding, 
          i * cellSize + padding, 
          cellSize, 
          cellSize
        );
        
        // Add value text
        ctx.fillStyle = normalized > 0.5 ? 'white' : 'black';
        ctx.font = '12px sans-serif';
        ctx.textAlign = 'center';
        ctx.textBaseline = 'middle';
        ctx.fillText(
          value.toFixed(2), 
          j * cellSize + padding + cellSize / 2, 
          i * cellSize + padding + cellSize / 2
        );
      }
    }
    
    // Draw axes labels
    ctx.fillStyle = 'black';
    ctx.font = '12px sans-serif';
    ctx.textAlign = 'center';
    
    for (let i = 0; i < numClasses; i++) {
      // Column labels (predicted)
      ctx.fillText(
        i.toString(), 
        i * cellSize + padding + cellSize / 2, 
        padding - 15
      );
      
      // Row labels (actual)
      ctx.textAlign = 'right';
      ctx.fillText(
        i.toString(), 
        padding - 10, 
        i * cellSize + padding + cellSize / 2
      );
      ctx.textAlign = 'center';
    }
    
    // Draw axes titles
    ctx.font = '14px sans-serif';
    ctx.fillText('Predicted', width / 2, padding - 30);
    
    // Rotate text for y-axis
    ctx.save();
    ctx.translate(padding - 30, height / 2);
    ctx.rotate(-Math.PI / 2);
    ctx.fillText('Actual', 0, 0);
    ctx.restore();
    
  }, [matrix]);
  
  if (!matrix || matrix.length === 0) {
    return (
      <div className="bg-white rounded-lg border border-gray-200 p-6 text-center shadow-sm">
        <p className="text-gray-500">No confusion matrix available</p>
      </div>
    );
  }
  
  return (
    <div className="bg-white rounded-lg border border-gray-200 p-4 shadow-sm">
      <h3 className="text-lg font-medium text-gray-800 mb-4">Confusion Matrix</h3>
      <div className="overflow-auto flex justify-center">
        <canvas 
          ref={canvasRef}
          className="bg-white"
        />
      </div>
      <p className="text-sm text-gray-500 mt-2 text-center">
        The confusion matrix shows how often the model predicts each digit correctly.
        Darker cells indicate more frequent predictions.
      </p>
    </div>
  );
};

export default ConfusionMatrix;